library(tidyverse)

ben_theme <- function(){
    theme_classic() +
        theme(panel.background = element_blank(),
              panel.grid.major = element_blank(),
              panel.grid.minor = element_blank(),
              axis.line = element_line(colour = "black"),
              panel.border = element_rect(colour = "black", fill = NA, size = 1),
              strip.background = element_blank(),
              legend.position = "bottom", legend.title = element_blank(),
              plot.title = element_text(hjust = 0.5))
}

## -------------------------------------------
## Full enumeration - 25 precinct, 2 districts
## -------------------------------------------
fn <- list.files("../data/qq_test/", pattern="*.csv")
nums <- parse_number(fn)
df_qq_rsg <- do.call("rbind", lapply(paste0("../data/qq_test/", fn[nums <= 200]),
                                 read_csv)) %>%
    mutate(
        parity = factor(parity, levels = c("No Equal Population Constraint",
                                           "20% Equal Population Constraint",
                                           "10% Equal Population Constraint"))
    ) %>% filter(alg == "RSG")
df_qq <- do.call("rbind", lapply(paste0("../data/qq_test/", fn[nums > 200]),
                                 read_csv)) %>%
    mutate(
        parity = factor(parity, levels = c("No Equal Population Constraint",
                                           "20% Equal Population Constraint",
                                           "10% Equal Population Constraint"))
    )
df_qq <- bind_rows(df_qq, df_qq_rsg)

## Create plot
df_qq %>%
    filter(alg == "RSG" |
           alg == "MCMC\n(Soft Constraint)" | 
           (alg == ("MCMC\n(Hard Constraint)") & parity == "No Equal Population Constraint")   ) %>%
    mutate(alg = case_when(alg == "RSG"~alg,
                           TRUE~"MCMC")) %>%
    ggplot(aes(sample = ks_pval, fill = alg, colour = alg, pch = alg)) + 
    stat_qq(distribution = stats::qunif) + 
    scale_color_manual(values = c("black", "red")) +
    scale_shape_manual(values = c(16, 3)) + 
    facet_grid(~ parity) +
    labs(x = "Theoretical Quantile", y = "Empirical Quantile") +
    ben_theme() +
    theme(legend.position = c(.07, .8)) +
    geom_abline(aes(intercept = 0, slope = 1), lty = "dashed") +
    ggsave("../paper/figs/qq_test.pdf", height = 3, width = 8)

#|(alg == ("MCMC\r\n(Hard Constraint)") & parity == "No Equal Population Constraint")            alg == "MCMC\r\n(Soft Constraint)"|

## ----------------------------------------------
## Sample enumeration - 100 precinct, 2 districts
## ----------------------------------------------
fn <- list.files("../data/qq_test_sample/", pattern="*.csv")
nums <- parse_number(fn)
df_qq <- do.call("rbind", lapply(paste0("../data/qq_test_sample/", fn), read_csv)) %>%
    mutate(
        parity = factor(parity, levels = c("No Equal Population Constraint",
                                           "5% Equal Population Constraint",
                                           "1% Equal Population Constraint"))
    ) 

## Create plot
ggplot(df_qq, aes(sample = ks_pval, fill = alg, colour = alg, pch = alg)) + 
    stat_qq(distribution = stats::qunif) + 
    scale_color_manual(values = c("black", "red")) +
    scale_shape_manual(values = c(16, 3)) + 
    facet_grid(~ parity) +
    labs(x = "Theoretical Quantile", y = "Empirical Quantile") +
    ben_theme() +
    theme(legend.position = c(.07, .85)) +
    geom_abline(aes(intercept = 0, slope = 1), lty = "dashed") +
    ggsave("../paper/figs/qq_test_sample.pdf", height = 3, width = 8)
